"""
Tushare 股票数据查询工具
统一的股票数据接口，供所有智能体使用
"""
import tushare as ts
from datetime import datetime, timedelta
import time
import random
from utils.retry_helper import retry_on_error
from urllib.error import HTTPError, URLError
import socket


# Tushare Token
TUSHARE_TOKEN = "52b44eb3a752a5c49b381a1308130e61e9beaf14070b4de07c02a441"


def init_tushare():
    """初始化 Tushare"""
    ts.set_token(TUSHARE_TOKEN)
    return ts.pro_api() if TUSHARE_TOKEN else None


def normalize_stock_code(symbol: str) -> str:
    """
    标准化股票代码
    支持多种输入格式：600000, 000001, 600000.SH, SH600000
    """
    symbol = symbol.upper().strip()
    symbol = symbol.replace('SH', '').replace('SZ', '').replace('.', '')

    if symbol.startswith('6') or symbol.startswith('5'):
        return f"{symbol}.SH"
    elif symbol.startswith('0') or symbol.startswith('3'):
        return f"{symbol}.SZ"
    else:
        return f"{symbol}.SH"


@retry_on_error(max_retries=3, delay=2.0, exceptions=(HTTPError, URLError, socket.timeout, ConnectionError))
def get_stock_price(symbol: str) -> str:
    """
    获取股票实时价格

    Args:
        symbol: 股票代码（如 600000, 000001）

    Returns:
        股票价格信息字符串
    """
    try:
        ts_code = normalize_stock_code(symbol)
        df = ts.get_realtime_quotes(symbol)

        if df.empty or df.iloc[0]['price'] == '0.00':
            return f"⚠️ 无法获取 {symbol} 的数据，可能是非交易时间或股票代码不存在"

        data = df.iloc[0]
        current_price = float(data['price'])
        pre_close = float(data['pre_close'])
        change = current_price - pre_close
        change_percent = (change / pre_close) * 100 if pre_close > 0 else 0

        result = f"""
【{symbol} ({ts_code}) 实时行情】
股票名称: {data['name']}
当前价格: ¥{current_price:.2f}
昨日收盘: ¥{pre_close:.2f}
涨跌额: ¥{change:+.2f}
涨跌幅: {change_percent:+.2f}%
今日开盘: ¥{float(data['open']):.2f}
最高价: ¥{float(data['high']):.2f}
最低价: ¥{float(data['low']):.2f}
成交量: {float(data['volume']):,.0f} 手
成交额: {float(data['amount']):,.2f} 万元
更新时间: {data['date']} {data['time']}
"""
        return result.strip()

    except Exception as e:
        return f"❌ 获取 {symbol} 价格失败: {str(e)}"


@retry_on_error(max_retries=3, delay=2.0, exceptions=(HTTPError, URLError, socket.timeout, ConnectionError))
def get_stock_info(symbol: str) -> str:
    """
    获取股票基本信息

    Args:
        symbol: 股票代码

    Returns:
        股票详细信息字符串
    """
    try:
        time.sleep(random.uniform(0.3, 0.8))
        ts_code = normalize_stock_code(symbol)
        df = ts.get_realtime_quotes(symbol)

        if df.empty:
            return f"⚠️ 无法获取 {symbol} 的信息"

        stock_name = df.iloc[0]['name']

        # 行业信息映射
        industry_map = {
            '600000': ('银行', '上海'),
            '601398': ('银行', '北京'),
            '601288': ('银行', '北京'),
            '600036': ('银行', '深圳'),
            '600519': ('食品饮料', '贵州'),
            '000858': ('食品饮料', '四川'),
            '002415': ('电子', '浙江'),
            '300750': ('电气设备', '福建'),
        }

        industry, area = industry_map.get(symbol, ('未知行业', '未知地区'))

        result = f"""
【{symbol} ({ts_code}) 公司信息】
股票名称: {stock_name}
所属行业: {industry}
所属地区: {area}
上市市场: {'上海证券交易所' if ts_code.endswith('.SH') else '深圳证券交易所'}

💡 提示：详细财务数据需要 Tushare 积分权限
"""
        return result.strip()

    except Exception as e:
        return f"❌ 获取 {symbol} 信息失败: {str(e)}"


def get_stock_history(symbol: str, days: int = 30) -> str:
    """
    获取股票历史K线数据

    Args:
        symbol: 股票代码
        days: 天数（默认30天）

    Returns:
        历史数据摘要字符串
    """
    try:
        time.sleep(random.uniform(0.3, 0.8))
        ts_code = normalize_stock_code(symbol)

        # 计算日期范围
        end_date = datetime.now()
        start_date = end_date - timedelta(days=days+10)  # 多取几天以确保有足够数据

        start_date_str = start_date.strftime('%Y%m%d')
        end_date_str = end_date.strftime('%Y%m%d')

        # 获取真实的日K数据
        pro = init_tushare()
        if not pro:
            return f"⚠️ Tushare 未初始化，无法获取历史数据"

        # 根据市场类型选择接口
        market = ts_code.split(".")[-1].upper()

        if market == "SH":
            if ts_code.startswith("000"):
                # 上证指数
                df = pro.index_daily(ts_code=ts_code, start_date=start_date_str, end_date=end_date_str)
            else:
                # A股股票
                df = pro.daily(ts_code=ts_code, start_date=start_date_str, end_date=end_date_str)
        elif market == "SZ":
            if ts_code.startswith("399"):
                # 深证指数
                df = pro.index_daily(ts_code=ts_code, start_date=start_date_str, end_date=end_date_str)
            else:
                # A股股票
                df = pro.daily(ts_code=ts_code, start_date=start_date_str, end_date=end_date_str)
        else:
            # 默认使用 daily 接口
            df = pro.daily(ts_code=ts_code, start_date=start_date_str, end_date=end_date_str)

        if df is None or df.empty:
            return f"⚠️ 无法获取 {symbol} 的历史数据"

        # 按日期排序（从旧到新）
        df = df.sort_values('trade_date')

        # 只取最近的 days 天
        df = df.tail(days)

        if len(df) == 0:
            return f"⚠️ 无法获取 {symbol} 的历史数据"

        # 计算统计数据
        period_high = df['high'].max()
        period_low = df['low'].min()
        avg_price = df['close'].mean()
        first_price = df.iloc[0]['close']
        last_price = df.iloc[-1]['close']
        change = last_price - first_price
        change_percent = (change / first_price) * 100 if first_price > 0 else 0

        # 格式化最近5天的数据
        recent_5 = df.tail(5)
        recent_data = []

        for idx, row in recent_5.iterrows():
            date_str = str(row['trade_date'])
            formatted_date = f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:]}"
            day_change = row['close'] - row['open']
            day_change_pct = (day_change / row['open'] * 100) if row['open'] > 0 else 0

            recent_data.append(
                f"  {formatted_date}: 开¥{row['open']:.2f} 收¥{row['close']:.2f} "
                f"高¥{row['high']:.2f} 低¥{row['low']:.2f} 量{row['vol']:,.0f}手 "
                f"涨跌{day_change_pct:+.2f}%"
            )

        result = f"""
【{symbol} ({ts_code}) 近{len(df)}天历史数据】
期间最高价: ¥{period_high:.2f}
期间最低价: ¥{period_low:.2f}
期间平均价: ¥{avg_price:.2f}
期初价格: ¥{first_price:.2f} ({df.iloc[0]['trade_date']})
最新价格: ¥{last_price:.2f} ({df.iloc[-1]['trade_date']})
期间涨跌: ¥{change:+.2f} ({change_percent:+.2f}%)
数据点数: {len(df)} 天

最近5天数据:
{chr(10).join(recent_data)}

✅ 数据来源：Tushare 真实K线数据
"""
        return result.strip()

    except Exception as e:
        error_msg = str(e)
        if "权限" in error_msg or "permission" in error_msg.lower():
            return f"""
⚠️ Tushare API 权限不足
无法访问 daily() 接口获取历史K线数据
需要积分权限才能使用此功能

详情: https://tushare.pro/document/1?doc_id=108
"""
        return f"❌ 获取历史数据失败: {error_msg}"


# 常用股票代码参考
POPULAR_STOCKS = {
    "银行股": ["600000", "601398", "601288", "600036", "601166"],
    "科技股": ["002415", "300750", "688981"],
    "消费股": ["600519", "000858", "600887", "000568"],
    "地产股": ["000002", "600048", "001979"]
}


if __name__ == "__main__":
    print("=" * 80)
    print("Tushare 股票数据工具测试")
    print("=" * 80)
    print()

    test_symbol = "600000"

    print("1. 测试获取股票价格")
    print(get_stock_price(test_symbol))
    print()

    time.sleep(1)

    print("2. 测试获取股票信息")
    print(get_stock_info(test_symbol))
    print()

    time.sleep(1)

    print("3. 测试获取历史数据")
    print(get_stock_history(test_symbol, 30))
    print()

    print("=" * 80)
    print("✅ 测试完成")
    print("=" * 80)
